import numpy as np
import matplotlib.pyplot as plt
import os
import math


def load_and_plot_npy(file_path, save_plot=False, points_per_plot=2000):
    """
    简单加载NPY文件并绘制波形图，每2000个点一张图

    Args:
        file_path: NPY文件路径
        save_plot: 是否保存图片
        points_per_plot: 每张图显示的数据点数
    """
    # 加载数据
    try:
        data = np.load(file_path)
        print(f"成功加载: {os.path.basename(file_path)}")
        print(f"数据形状: {data.shape}")
        print(f"数据类型: {data.dtype}")
        print(f"数据总数: {len(data)}")
    except Exception as e:
        print(f"加载文件失败: {e}")
        return

    # 计算需要绘制的图数量
    total_points = len(data)
    num_plots = math.ceil(total_points / points_per_plot)
    print(f"每张图显示点数: {points_per_plot}")
    print(f"总共需要绘制: {num_plots} 张图")

    # 逐张绘制图片
    for i in range(num_plots):
        # 计算当前图的起始和结束索引
        start_idx = i * points_per_plot
        end_idx = min((i + 1) * points_per_plot, total_points)

        # 提取当前段的数据
        segment_data = data[start_idx:end_idx]

        # 创建图形
        plt.figure(figsize=(12, 6))

        # 绘制波形
        plt.plot(segment_data, linewidth=1)
        plt.title(f'波形图 - {os.path.basename(file_path)} (第{i + 1}/{num_plots}段)')
        plt.xlabel('采样点 (相对位置)')
        plt.ylabel('数值')
        plt.grid(True, alpha=0.3)

        # 添加统计信息
        stats_text = f'数据点数: {len(segment_data)}\n起始索引: {start_idx}\n结束索引: {end_idx - 1}\n最小值: {np.min(segment_data):.2f}\n最大值: {np.max(segment_data):.2f}\n平均值: {np.mean(segment_data):.2f}'
        plt.text(0.02, 0.98, stats_text, transform=plt.gca().transAxes,
                 verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))

        plt.tight_layout()

        # 保存图片
        if save_plot:
            output_path = file_path.replace('.npy', f'_plot_part_{i + 1}.png')
            plt.savefig(output_path, dpi=150, bbox_inches='tight')
            print(f"图片已保存: {output_path}")

        # 显示图片
        plt.show()

        print(f"已完成第 {i + 1}/{num_plots} 张图绘制")


def main():
    # 直接指定文件路径
    file_path = R"D:\2025_CXYZP008\20250910__brain_slam\sync_sample.npy"  # 修改为你的NPY文件路径

    if os.path.exists(file_path):
        load_and_plot_npy(file_path, save_plot=True, points_per_plot=500000)
    else:
        print(f"文件不存在: {file_path}")
        print("请修改脚本中的 file_path 变量为你的NPY文件路径")


if __name__ == "__main__":
    main()